from flask_restx import Resource, fields
from flask import request, Response, stream_with_context
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import PromptTemplate
import json
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))

# 导入load_prompt函数
from langchain_core.prompts import load_prompt

# 定义数据模型用于请求参数验证
chat_model_input = None
clear_history_model = None
error_model = None

# 全局变量占位符
chat_model = None
session_histories = None
chat_ns = None
prompt_template = None

def init_chat_api(api_chat_ns, model, histories):
    """初始化聊天API"""
    global chat_model, session_histories, chat_ns, prompt_template
    global chat_model_input, clear_history_model, error_model
    
    # 设置全局变量
    chat_model = model
    session_histories = histories
    chat_ns = api_chat_ns

    # 定义数据模型用于请求参数验证
    chat_model_input = chat_ns.model('ChatInput', {
        'message': fields.String(required=True, description='用户输入的消息'),
        'session_id': fields.String(required=False, default='default', description='会话ID')
    })

    clear_history_model = chat_ns.model('ClearHistoryInput', {
        'session_id': fields.String(required=False, default='default', description='会话ID')
    })

    # 错误响应模型
    error_model = chat_ns.model('Error', {
        'error': fields.String(description='错误信息')
    })
    
    # 加载提示词模板
    try:
        # 获取项目根目录
        project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
        prompt_path = os.path.join(project_root, "templates", "prompt", "prompt.json")
        prompt_template = load_prompt(prompt_path, encoding="utf-8")
    except Exception as e:
        # 如果加载失败，使用默认模板
        prompt_template = PromptTemplate.from_template("请回答以下问题: {message}")

    # 添加路由和方法
    @chat_ns.route('')
    class Chat(Resource):
        @chat_ns.expect(chat_model_input)
        @chat_ns.produces(['text/event-stream'])
        @chat_ns.response(200, 'Success')
        @chat_ns.response(400, 'Bad Request', error_model)
        @chat_ns.response(500, 'Internal Server Error', error_model)
        def post(self):
            """处理聊天请求"""
            try:
                data = request.json
                user_input = data.get('message')
                session_id = data.get('session_id', 'default')
                
                if not user_input:
                    return {'error': '消息不能为空'}, 400

                # 使用提示词模板格式化用户输入
                try:
                    # 根据模板变量动态处理
                    formatted_input = prompt_template.format(
                        name="你" if "name" in prompt_template.input_variables else "用户",
                        what=user_input if "what" in prompt_template.input_variables else "有趣",
                        message=user_input if "message" in prompt_template.input_variables else user_input
                    )
                except Exception:
                    # 如果模板格式化失败，直接使用原始输入
                    formatted_input = user_input

                # 获取或创建会话历史
                if session_id not in session_histories:
                    session_histories[session_id] = []

                # 添加用户消息到历史
                session_histories[session_id].append(HumanMessage(content=formatted_input))

                # 构建消息历史
                messages = session_histories[session_id][-10:]  # 限制历史消息数量

                # 调用模型并流式输出
                def generate():
                    try:
                        full_response = ""
                        for chunk in chat_model.stream(messages):
                            full_response += chunk.content
                            # 以 SSE 格式发送数据
                            yield f"data: {json.dumps({'content': chunk.content})}\n\n"

                        # 将完整回复添加到历史记录
                        session_histories[session_id].append(AIMessage(content=full_response))

                        # 发送结束信号
                        yield "data: [DONE]\n\n"
                    except Exception as e:
                        yield f"data: {json.dumps({'error': str(e)})}\n\n"

                return Response(stream_with_context(generate()), content_type='text/event-stream')

            except Exception as e:
                return {'error': str(e)}, 500

    @chat_ns.route('/history')
    class ClearHistory(Resource):
        @chat_ns.expect(clear_history_model)
        @chat_ns.response(200, 'Success')
        @chat_ns.response(500, 'Internal Server Error', error_model)
        def post(self):
            """清除会话历史"""
            try:
                data = request.json
                session_id = data.get('session_id', 'default')
                if session_id in session_histories:
                    session_histories[session_id] = []
                return {'message': '历史记录已清除'}
            except Exception as e:
                return {'error': str(e)}, 500